#include <iostream>
#include <fstream>
#include <vector>
#include <cstdint>
#include <string>
#include <nlohmann/json.hpp>
#include "pointcloud_tools.h"
#include "point_preprocess.h"
#include "point_common.h"
#include "model_builder.h"

using json = nlohmann::json;
using namespace pointcloud_tools;
using namespace PointPreprocess;

namespace model_builder
{
    /*
    Our model consists of two types of OP:
    1. SubmConv Layer
    2. FC Layer

    TODO:
    1. Profiler
    2. Export results for debugging
    3. Optimize performance
    4. Support other OP types
    */

    //---------------------------SubmConv Operation---------------------------

    void SubmConvLayerConfig::run(ModelBuilder &builder, Pointcloud_utils &ptu)
    {
        // 1. initialize weight and input
        ptu.wt_init();
        ptu.compute_init();

        // 2. run sconv
        ptu.run();

        // 3. merge result
        builder.merged_res.clear();
        // int och_loop_num = get_used_core_num(ptu.layer_config.cfg_core_enable) - 1;
        /*if (! bias_mode)
            merge_RESULT_feat(ptu.block_config, builder.input_pt_num, cfg_och_loop_num, builder.pt_mapping, builder.valid_filter, builder.merged_res);
        else*/
        // std::cout << "layer relu enable: " << relu_enable << std::endl;
        merge_RESULT_feat(ptu.block_config, builder.input_pt_num, cfg_och_loop_num, builder.pt_mapping, builder.valid_filter, builder.merged_res, bias_cores, relu_enable);
    }

    void SubmConvLayerConfig::run(ModelBuilder &builder, Pointcloud_online_utils &ptu)
    {
        // 1. initialize weight
        ptu.wt_init();

        // 2. initialize instructions
        int cnn_inst_index = ptu.inst_gen_block(builder.dummy_bcfg);

        // 3. run sconv
        ptu.run_block(cnn_inst_index);

        // 4. merge result
        builder.merged_res.clear();
        // int och_loop_num = get_used_core_num(ptu.layer_config.cfg_core_enable) - 1;
        // std::cout << "layer relu enable: " << relu_enable << std::endl;
        merge_RESULT_feat(ptu.block_config, builder.input_pt_num, cfg_och_loop_num, builder.pt_mapping, builder.valid_filter, builder.merged_res, bias_cores, relu_enable);
    }

    void SubmConvLayerConfig::UpdateInput(ModelBuilder &builder, Pointcloud_utils &ptu)
    {
        // 1. update ptu
        ptu.layer_config = layer_config;
        ptu.layer_config.BASE_WT_ADDR = layer_config.BASE_WT_ADDR;

        // 2. update input block
        uint64_t DDR_IN_ADDR = ptu.block_config[0].BASE_IN_ADDR;
        uint64_t DDR_RS_ADDR = ptu.block_config.back().BASE_RS_ADDR;
        // int och_loop_num = get_used_core_num(layer_config.cfg_core_enable) - 1;
        for (int block_idx = 0; block_idx < ptu.block_config.size(); block_idx++)
        {
            ptu.block_config[block_idx].in_data.clear();
            int pt_num = ptu.block_config[block_idx].cfg_cam_valid_num;
            auto pt_mapping_tmp = builder.pt_mapping[block_idx];
            ptu.block_config[block_idx].in_data.reserve(pt_num * (ptu.layer_config.cfg_inch_loop_num + 1));
            for (int j = 0; j < ptu.layer_config.cfg_inch_loop_num + 1; j++)
            {
                for (int k = 0; k < pt_num; k++)
                {
                    int idx_tmp = pt_mapping_tmp[k] + j * builder.input_pt_num;
                    ptu.block_config[block_idx].in_data.push_back(builder.merged_res[idx_tmp]);
                }
            }
            ptu.block_config[block_idx].BASE_IN_ADDR = DDR_IN_ADDR;
            ptu.block_config[block_idx].BASE_RS_ADDR = DDR_RS_ADDR;
            DDR_IN_ADDR += pt_num * sizeof(uint64_t) * (1 + layer_config.cfg_inch_loop_num);
            DDR_RS_ADDR += pt_num * sizeof(uint64_t) * (1 + cfg_och_loop_num);
        }
    }

    void SubmConvLayerConfig::UpdateInput(ModelBuilder &builder, Pointcloud_online_utils &ptu)
    {
        // 1. update ptu
        ptu.layer_config = layer_config;

        // 2. update input block
        uint64_t DDR_IN_ADDR = ptu.block_config[0].BASE_IN_ADDR;
        uint64_t DDR_RS_ADDR = ptu.block_config.back().BASE_RS_ADDR;

        for (int block_idx = 0; block_idx < ptu.block_config.size(); block_idx++)
        {
            ptu.block_config[block_idx].in_data.clear();
            int pt_num = ptu.block_config[block_idx].cfg_cam_valid_num;
            auto pt_mapping_tmp = builder.pt_mapping[block_idx];
            ptu.block_config[block_idx].in_data.reserve(pt_num * (ptu.layer_config.cfg_inch_loop_num + 1));
            for (int j = 0; j < ptu.layer_config.cfg_inch_loop_num + 1; j++)
            {
                for (int k = 0; k < pt_num; k++)
                {
                    int idx_tmp = pt_mapping_tmp[k] + j * builder.input_pt_num;
                    ptu.block_config[block_idx].in_data.push_back(builder.merged_res[idx_tmp]);
                }
            }
            ptu.block_config[block_idx].BASE_IN_ADDR = DDR_IN_ADDR;
            ptu.block_config[block_idx].BASE_RS_ADDR = DDR_RS_ADDR;
        }
    }

    //---------------------------FC Operation---------------------------

    void FCLayerConfig::fc_gemm(ModelBuilder &builder)
    {
        // // UNFINISHED
        // Data64to8 packed_in;
        // std::vector<uint64_t> res_tmp;
        // res_tmp.resize(builder.input_pt_num * FC_OCH_LOOP_NUM, 0);

        // for (int m = 0; m < builder.input_pt_num; m++)
        // {
        //     // matrix mul: [builder.input_pt_num, ICH] * [ICH, OCH] = [builder.input_pt_num, OCH]
        //     for (int n = 0; n < FC_OCH_LOOP_NUM; n++)
        //     {
        //         for (int k = 0; k < FC_ICH_LOOP_NUM; k++)
        //         {
        //             packed_in.in64 = builder.merged_res[builder.input_pt_num * k + m];
        //             for (int i = 0; i < 8; i++)
        //             {
        //                 if (i < FC_ICH_LOOP_NUM % 8)
        //                 {
        //                     res_tmp[m * FC_OCH_LOOP_NUM + n] += packed_in.out8[i] * weights[n * FC_ICH_LOOP_NUM + k * 8 + i]
        //                 }
        //                 else
        //                 {
        //                     res_tmp[m * FC_OCH_LOOP_NUM + n] += 0;
        //                 }
        //             }
        //         }
        //     }
        // }

        // // res_tmp copied to merged_res
        // builder.merged_res.clear();
        // builder.merged_res.resize(builder.input_pt_num * FC_OCH_LOOP_NUM, 0);
        // for (int i = 0; i < builder.input_pt_num * FC_OCH_LOOP_NUM; i++)
        // {
        //     builder.merged_res[i] = res_tmp[i];
        // }
        return ;
    }

    void FCLayerConfig::run(ModelBuilder &builder, Pointcloud_utils &ptu)
    {
        fc_gemm(builder);
    }

    void FCLayerConfig::run(ModelBuilder &builder, Pointcloud_online_utils &ptu)
    {
        fc_gemm(builder);
    }

    void FCLayerConfig::UpdateInput(ModelBuilder &builder, Pointcloud_utils &ptu)
    {
        return;
    }

    void FCLayerConfig::UpdateInput(ModelBuilder &builder, Pointcloud_online_utils &ptu)
    {
        return;
    }

    // ---------------------------ModelBuilder Implementation---------------------------

    void ModelBuilder::LoadModelConfig(const std::string &model_config_path)
    {
        std::cout << "Loading model configuration from: " << model_config_path << std::endl;
        std::ifstream ifs(model_config_path);
        if (!ifs.is_open())
        {
            throw std::runtime_error("Cannot open JSON config: " + model_config_path);
        }

        json j;
        ifs >> j;

        // 1. Load the input point cloud data
        size_t dlen = readBinToVec(input.pt_data, j.at("input").at("input_coords").get<std::string>());
        assert(dlen > 0);
        dlen = readBinToVec(input.in_data, j.at("input").at("input_feats").get<std::string>());
        assert(dlen > 0);
        for (int i = 0; i < 6; ++i)
        {
            input.spatial_shape[i] = j.at("input").at("spatial_shape")[i].get<uint16_t>();
        }
        input.cfg_cam_valid_num = input.pt_data.size();
        input.cfg_total_point_num = input.pt_data.size();
        input_pt_num = input.cfg_total_point_num;

        // 2. Load layer configurations
        for (auto &layer : j.at("layers"))
        {
            std::string type = layer.at("type").get<std::string>();
            int layer_id = layer.at("layer_id").get<int>();

            if (type == "SUBMCONV")
            {
                auto &c = layer.at("config");
                LayerConfig cfg;
                cfg.layer_id = layer_id;
                cfg.cfg_inch_loop_num = c.at("cfg_inch_loop_num").get<uint8_t>();
                cfg.cfg_start_pt_addr = c.at("cfg_start_pt_addr").get<uint32_t>();
                cfg.cfg_shift_num = c.at("cfg_shift_num").get<uint8_t>();
                cfg.cfg_pe_relu_enable = c.at("cfg_pe_relu_enable").get<uint8_t>();
                for (int i = 0; i < 3; ++i)
                {
                    cfg.cfg_cam_search_range[i] = c.at("cfg_cam_search_range")[i].get<uint8_t>();
                    cfg.cfg_cnn_size[i] = c.at("cfg_cnn_size")[i].get<uint8_t>();
                }
                cfg.cfg_bias_enable = c.at("cfg_bias_enable").get<uint8_t>();
                cfg.cfg_cam_enable = c.at("cfg_cam_enable").get<uint8_t>();
                cfg.cfg_core_enable = c.at("cfg_core_enable").get<uint8_t>();
                for (int i = 0; i < 8; ++i)
                {
                    cfg.bias_core0[i] = c.at("bias_core0")[i].get<int8_t>();
                    cfg.bias_core1[i] = c.at("bias_core1")[i].get<int8_t>();
                    cfg.bias_core2[i] = c.at("bias_core2")[i].get<int8_t>();
                    cfg.bias_core3[i] = c.at("bias_core3")[i].get<int8_t>();
                }
                cfg.BASE_WT_ADDR = c.at("BASE_WT_ADDR").get<uint64_t>();
                cfg.DMA_WT_NUM = c.at("DMA_WT_NUM").get<uint16_t>();

                std::string weight_path = c.value("weight_path", std::string{});
                if (!weight_path.empty())
                {
                    dlen = readBinToVec(cfg.wt_data, weight_path);
                }

                addSubmConvLayer(cfg);
            }
            else if (type == "FC")
            {
                uint32_t ICH = layer.at("ICH").get<uint32_t>();
                uint32_t OCH = layer.at("OCH").get<uint32_t>();
                auto bias = layer.at("bias").get<std::vector<float>>();

                std::vector<float> weights;
                std::string weight_path = layer.value("weight_path", std::string{});
                if (!weight_path.empty())
                {
                    dlen = readBinToVec(weights, weight_path);
                }
                else
                {
                    weights.resize(static_cast<size_t>(ICH) * OCH);
                }

                addFCLayer(weights, bias, ICH, OCH);
            }
            else
            {
                std::cerr << "Unknown layer type: " << type << "\n";
            }
        }
        std::cout << "Model configuration loaded successfully!\n";
    }

    void ModelBuilder::PreprocessInput()
    {
        assert(layers[0]->type == LayerType::SUBMCONV);
        assert(!layers.empty());
        auto *input_layer = static_cast<SubmConvLayerConfig *>(layers[0]);
        octree_partition(input, 1024,
                         pt_mapping, bcfg_partition, valid_filter,
                         input_layer->layer_config.cfg_cam_search_range,
                         input.spatial_shape,
                         input_layer->layer_config.cfg_inch_loop_num,
                         input_layer->cfg_och_loop_num,
                         offline_mode);
    }

    void ModelBuilder::RunOfflineModel()
    {
        printf("Simulation Mode is ON!\n");

        // 1. configure the tools
        Pointcloud_utils ptu;
        ptu.ddrSize = 8 * 8 * 1024 * 1024 * sizeof(int); // 256MB
        bool is_alloc = ptu.Model_DDR_alloc("DDR_32MB.bin");
        assert(is_alloc);
        ptu.DDR_BASE_ADDR = DDR_BASE_ADDR;

        // 2. preprocess the input point cloud data
        PreprocessInput();
        for (int i = 0; i < bcfg_partition.size(); i++)
        {
            ptu.block_config.push_back(bcfg_partition[i]);
        }

        // 3. run each layer
        auto *L = layers[0];
        assert(L->type == LayerType::SUBMCONV);
        auto *input_layer = static_cast<SubmConvLayerConfig *>(L);
        ptu.layer_config = input_layer->layer_config;
        L->run(*this, ptu);
        std::string layer0_output_file = "layer0_output.txt";
        export_merged_results(merged_res, input_pt_num, input_layer->cfg_och_loop_num, layer0_output_file);

        for (size_t idx = 1; idx < layers.size(); ++idx)
        {
            std::cout << "Running layer " << idx + 1 << "...\n";
            auto *L = layers[idx];
            L->UpdateInput(*this, ptu);
            L->run(*this, ptu);
            if (L->type == LayerType::SUBMCONV)
            {
                std::string layer_output_file = "layer" + std::to_string(idx) + "_output.txt";
                int och_loop_num = get_used_core_num(ptu.layer_config.cfg_core_enable) - 1;
                export_merged_results(merged_res, input_pt_num, och_loop_num, layer_output_file);
            }
        }

        std::cout << "Offline Inference Complete !" << std::endl;
    }

    void ModelBuilder::RunOnlineModel()
    {
        printf("Onchip Mode is ON!\n");

        // 1. configure the tools
        Pointcloud_online_utils ptu;
        ptu.pmem_mmap_init();
        ptu.DDR_BASE_ADDR = DDR_BASE_ADDR;

        // 2. preprocess the input point cloud data
        PreprocessInput();
        for (int i = 0; i < bcfg_partition.size(); i++)
        {
            ptu.block_config.push_back(bcfg_partition[i]);
        }

        // 3. run each layer
        auto *L = layers[0];
        assert(L->type == LayerType::SUBMCONV);
        auto *input_layer = static_cast<SubmConvLayerConfig *>(L);
        ptu.layer_config = input_layer->layer_config;
        L->run(*this, ptu);
        // std::string layer0_output_file = "layer0_output.txt";
        // export_merged_results(merged_res, input_pt_num, input_layer->cfg_och_loop_num, layer0_output_file);

        for (size_t idx = 1; idx < layers.size(); ++idx)
        {
            auto *L = layers[idx];
            L->UpdateInput(*this, ptu);
            L->run(*this, ptu);
            // if (idx == layers.size() - 1 && L->type == LayerType::SUBMCONV)
            // {
            //     std::string layer_output_file = "layer" + std::to_string(idx) + "_output.txt";
            //     int och_loop_num = get_used_core_num(ptu.layer_config.cfg_core_enable) - 1;
            //     export_merged_results(merged_res, input_pt_num, och_loop_num, layer_output_file);
            // }
        }

        std::cout << "Online Inference Complete !" << std::endl;
    }

    void ModelBuilder::RunModel()
    {
        if (offline_mode)
        {
            RunOfflineModel();
        }
        else
        {
            // profiling the run time
            auto start = std::chrono::high_resolution_clock::now();
            RunOnlineModel();
            auto end = std::chrono::high_resolution_clock::now();
            std::chrono::duration<double, std::milli> duration = end - start;
            std::cout << "Online Inference Time: " << duration.count() << " ms" << std::endl;
        }
    }
}